{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# heatmap visualization\n",
    "import sys\n",
    "\n",
    "import pickle\n",
    "import re\n",
    "import os\n",
    "sys.path.append(\"..\")\n",
    "\n",
    "import random\n",
    "import numpy as np\n",
    "import torch\n",
    "from random import shuffle\n",
    "import argparse\n",
    "import pickle\n",
    "\n",
    "import collections\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "\n",
    "from model.QACGBERT import *\n",
    "\n",
    "from sklearn.metrics import f1_score, accuracy_score, roc_auc_score\n",
    "\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "from torch.utils.data.distributed import DistributedSampler\n",
    "from torch.utils.data.sampler import RandomSampler, SequentialSampler\n",
    "from tqdm import tqdm, trange\n",
    "\n",
    "from util.optimization import BERTAdam\n",
    "from util.processor import (Sentihood_NLI_M_Processor,\n",
    "                            Semeval_NLI_M_Processor)\n",
    "\n",
    "from util.tokenization import *\n",
    "\n",
    "from util.evaluation import *\n",
    "\n",
    "from util.processor import *\n",
    "from run_classifier import *\n",
    "from util.train_helper import *\n",
    "\n",
    "import seaborn as sns; sns.set()\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "plt.rcParams[\"font.family\"] = \"Times New Roman\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Set-ups"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note that this notebook only supports single GPU evaluation\n",
    "# which is sufficient for most of tasks by using lower batch size.\n",
    "IS_CUDA = False\n",
    "if IS_CUDA:\n",
    "    CUDA_DEVICE = \"cuda:2\"\n",
    "    device = torch.device(CUDA_DEVICE)\n",
    "    n_gpu = torch.cuda.device_count()\n",
    "    logger.info(\"device %s in total n_gpu %d distributed training\", device, n_gpu)\n",
    "else:\n",
    "    # bad luck, we are on CPU now!\n",
    "    logger.info(\"gpu is out of the picture, let us use CPU\")\n",
    "    device = torch.device(\"cpu\")\n",
    "    \n",
    "model, optimizer, tokenizer = \\\n",
    "    getModelOptimizerTokenizer(\"QACGBERT\", \n",
    "                               vocab_file=\"../../models/BERT-Google/vocab.txt\",\n",
    "                               bert_config_file=\"../../models/BERT-Google/bert_config.json\",\n",
    "                               init_checkpoint=\"../../results/sentihood/QACGBERT-1/best_checkpoint.bin\",\n",
    "                               label_list=['None', 'Positive', 'Negative'],\n",
    "                               num_train_steps=1,\n",
    "                               learning_rate=2e-5,\n",
    "                               base_learning_rate=2e-5,\n",
    "                               warmup_proportion=0.1,\n",
    "                               init_lrp=True)\n",
    "import argparse\n",
    "args = argparse.ArgumentParser()\n",
    "args.task_name = \"sentihood_NLI_M\"\n",
    "vocab_data_dir = \"../../models/BERT-Google/vocab.txt\"\n",
    "\n",
    "def heatmap_viz(token_grad, vmin=0, vmax=1, cmap=\"Blues\"):\n",
    "    scores = [tu[1] for tu in token_grad]\n",
    "    tokens = [tu[0] for tu in token_grad]\n",
    "    fig, ax = plt.subplots(figsize=(15,1))\n",
    "    ax = sns.heatmap([scores], cmap=cmap, xticklabels=tokens, yticklabels=False,\n",
    "                     cbar_kws=dict(shrink=1, aspect=4, ), linewidths=0.8)\n",
    "    ax.set_xticklabels(tokens, size = 18)\n",
    "    cbar = ax.collections[0].colorbar\n",
    "    # here set the labelsize by 20\n",
    "    cbar.ax.tick_params(labelsize=20)\n",
    "    plt.show()\n",
    "    \n",
    "def inverse_mapping(vocab_dict):\n",
    "    inverse_vocab_dict = {}\n",
    "    for k, v in vocab_dict.items():\n",
    "        inverse_vocab_dict[v] = k\n",
    "    return inverse_vocab_dict\n",
    "\n",
    "def translate(token_ids, vocab):\n",
    "    tokens = []\n",
    "    for _id in token_ids.tolist():\n",
    "        tokens.append(vocab[_id])\n",
    "    return tokens\n",
    "\n",
    "vocab = inverse_mapping(load_vocab(vocab_data_dir, pretrain=False))\n",
    "\n",
    "context_id_map_sentihood = ['location - 1 - general',\n",
    "                            'location - 1 - price',\n",
    "                            'location - 1 - safety',\n",
    "                            'location - 1 - transit location',\n",
    "                            'location - 2 - general',\n",
    "                            'location - 2 - price',\n",
    "                            'location - 2 - safety',\n",
    "                            'location - 2 - transit location']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_examples = Sentihood_NLI_M_Processor().get_test_examples(\"../../datasets/sentihood\")\n",
    "test_features = convert_examples_to_features(\n",
    "    test_examples, ['None', 'Positive', 'Negative'], 512,\n",
    "    tokenizer, 1, True, args=args)\n",
    "\n",
    "all_input_ids = torch.tensor([f.input_ids for f in test_features], dtype=torch.long)\n",
    "all_input_mask = torch.tensor([f.input_mask for f in test_features], dtype=torch.long)\n",
    "all_segment_ids = torch.tensor([f.segment_ids for f in test_features], dtype=torch.long)\n",
    "all_label_ids = torch.tensor([f.label_id for f in test_features], dtype=torch.long)\n",
    "all_seq_len = torch.tensor([[f.seq_len] for f in test_features], dtype=torch.long)\n",
    "all_context_ids = torch.tensor([f.context_ids for f in test_features], dtype=torch.long)\n",
    "\n",
    "test_data = TensorDataset(all_input_ids, all_input_mask, all_segment_ids,\n",
    "                            all_label_ids, all_seq_len, all_context_ids)\n",
    "test_sampler = RandomSampler(test_data)\n",
    "test_dataloader = DataLoader(test_data, batch_size=1, sampler=test_sampler)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Single Sample Visualizations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()\n",
    "rand_sampler = torch.utils.data.RandomSampler(test_data, num_samples=1, replacement=True)\n",
    "train_sampler = torch.utils.data.DataLoader(test_data, batch_size=1, sampler=rand_sampler)\n",
    "\n",
    "for input_ids, input_mask, segment_ids, label_ids, seq_lens, context_ids in train_sampler:\n",
    "    max_seq_lens = max(seq_lens)[0]\n",
    "    input_ids = input_ids[:,:max_seq_lens]\n",
    "    input_mask = input_mask[:,:max_seq_lens]\n",
    "    segment_ids = segment_ids[:,:max_seq_lens]\n",
    "    tmp_test_loss, logits, all_new_attention_probs, all_attention_probs, all_quasi_attention_prob, _ = \\\n",
    "        model(input_ids, segment_ids, input_mask, seq_lens,\n",
    "                device=torch.device(\"cpu\"), labels=label_ids,\n",
    "                context_ids=context_ids)\n",
    "\n",
    "    # backing out using gradients\n",
    "    logits = F.softmax(logits, dim=-1)\n",
    "    sensitivity_class = 1\n",
    "    sensitivity_scores = torch.zeros(logits.shape)\n",
    "    sensitivity_scores[:,sensitivity_class] = 1.0\n",
    "    sensitivity_scores = model.backward_gradient(sensitivity_scores)\n",
    "    sensitivity_scores_norm = torch.norm(sensitivity_scores, dim=-1) * torch.norm(sensitivity_scores, dim=-1)\n",
    "    max_sensitivity_scores_norm = torch.max(sensitivity_scores_norm)\n",
    "    sensitivity_scores_norm = sensitivity_scores_norm / max_sensitivity_scores_norm\n",
    "    sensitivity_scores_norm = sensitivity_scores_norm.squeeze(dim=0).tolist()\n",
    "    print(\"*** pred logits ***\")\n",
    "    print(logits)\n",
    "    print(\"*** true label ***\")\n",
    "    print(label_ids[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"*** context ***\")\n",
    "print(context_id_map_sentihood[context_ids[0]])\n",
    "tokens = translate(input_ids[0], vocab)\n",
    "# hacky way to pick a head at random\n",
    "print(\"*** gradient plots ***\")\n",
    "heatmap_viz(list(zip(tokens, sensitivity_scores_norm)), cmap=\"Blues\", vmin=0, vmax=1)\n",
    "\n",
    "l = 0\n",
    "h = random.randint(0, 11)\n",
    "print(\"*** a_h plots ***\")\n",
    "a_h = all_new_attention_probs[l].squeeze(dim=0)[h].sum(dim=0).data\n",
    "a_h = a_h / torch.max(torch.abs(a_h))\n",
    "heatmap_viz(list(zip(tokens, a_h)), cmap=\"Blues\", vmin=-0.5, vmax=0.5)\n",
    "\n",
    "# print(\"*** a_self plots ***\")\n",
    "# a_self = all_attention_probs[l].squeeze(dim=0)[h].sum(dim=0).data\n",
    "# a_self = a_self / torch.max(torch.abs(a_self))\n",
    "# heatmap_viz(list(zip(tokens, a_self)), cmap=\"Blues\", vmin=0, vmax=1)\n",
    "\n",
    "print(\"*** a_quasi plots ***\")\n",
    "a_quasi = all_quasi_attention_prob[l].squeeze(dim=0)[h].sum(dim=0).data\n",
    "a_quasi = (a_quasi / torch.max(torch.abs(a_quasi)))\n",
    "heatmap_viz(list(zip(tokens, a_quasi)), cmap=\"Reds_r\", vmin=-0.5, vmax=0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Histgrams of Weight Matrices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()\n",
    "rand_sampler = torch.utils.data.RandomSampler(test_data, num_samples=5000, replacement=True)\n",
    "train_sampler = torch.utils.data.DataLoader(test_data, batch_size=1, sampler=rand_sampler)\n",
    "\n",
    "samples_lambda_context = []\n",
    "attention_probs = []\n",
    "quasi_attention_scores = []\n",
    "new_attention_probs = []\n",
    "count = 0\n",
    "\n",
    "for input_ids, input_mask, segment_ids, label_ids, seq_lens, context_ids in train_sampler:\n",
    "    max_seq_lens = max(seq_lens)[0]\n",
    "    input_ids = input_ids[:,:max_seq_lens]\n",
    "    input_mask = input_mask[:,:max_seq_lens]\n",
    "    segment_ids = segment_ids[:,:max_seq_lens]\n",
    "    tmp_test_loss, logits, all_new_attention_probs, all_attention_probs, all_quasi_attention_prob, all_lambda_context = \\\n",
    "        model(input_ids, segment_ids, input_mask, seq_lens,\n",
    "                device=torch.device(\"cpu\"), labels=label_ids,\n",
    "                context_ids=context_ids)\n",
    "    if label_ids.tolist()[0] != 0:\n",
    "        for i in range(12):\n",
    "            samples_lambda_context.extend(all_lambda_context[i].flatten().tolist())\n",
    "            attention_probs.extend(all_attention_probs[i].flatten().tolist())\n",
    "            quasi_attention_scores.extend(all_quasi_attention_prob[i].flatten().tolist())\n",
    "            new_attention_probs.extend(all_new_attention_probs[i].flatten().tolist())\n",
    "        count += 1\n",
    "        if count == 200:\n",
    "            break\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_histgram(array_in,  facecolor='g', xl=-1, xh=1):\n",
    "    import matplotlib as mpl\n",
    "    mpl.style.use(\"default\")\n",
    "    font = {'family' : 'Times New Roman',\n",
    "            'size'   : 30}\n",
    "    plt.rc('font', **font)\n",
    "    \n",
    "    fig = plt.figure(figsize=(6,6))\n",
    "    ax = fig.add_subplot(111)\n",
    "    g = ax.hist(array_in, bins=50, facecolor=facecolor)\n",
    "    plt.grid(True)\n",
    "    plt.grid(color='black', linestyle='-.')\n",
    "    import matplotlib.ticker as mtick\n",
    "    ax.yaxis.set_major_formatter(mtick.FormatStrFormatter('%.2e'))\n",
    "    ax.set_yscale('log')\n",
    "    plt.tight_layout()\n",
    "    plt.xlim(xl, xh)\n",
    "    plt.xticks([-1.5, -0.75, 0, 0.75, 1.5], [\"-1.5\", \"-0.75\", \"0.0\", \"0.75\", \"1.5\"], fontsize=30)\n",
    "    # plt.xticks([-0.5, -0.25, 0, 0.25, 0.5], [\"-0.50\", \"-0.25\", \"0.0\", \"0.25\", \"0.5\"], fontsize=30)\n",
    "    # plt.xticks([0, 0.25, 0.5, 0.75, 1], [\"0.0\", \"0.25\", \"0.50\", \"0.75\", \"1.0\"], fontsize=20)\n",
    "    # plt.xticks([0, 0.25, 0.5, 0.75, 1], [\"0.0\", \"0.25\", \"0.50\", \"0.75\", \"1.0\"], fontsize=20)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_histgram(new_attention_probs, facecolor=\"g\", xl=-1, xh=1)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
